package com.ruoyi.common.ssh;

import com.jcraft.jsch.*;
import com.ruoyi.common.utils.StringUtils;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * SSH连接管理器
 */
@Component
public class SSHConnectionManager {

    // 用于存储多个SSH连接对象，以主机名和端口号作为键
    private Map<String, Session> connectionMap = new HashMap<String, Session>();

    private Map<String, Channel> channelMap = new HashMap<String, Channel>();

    /**
     * 创建SSH连接的方法
     * @param host
     * @param port
     * @param username
     * @param password
     * @throws Exception
     */
    public boolean createConnection(String host, int port, String username, String password) throws Exception {
        JSch jsch = new JSch();
        Session session = jsch.getSession(username, host, port);
        session.setPassword(password);

        // 设置连接属性，这里为了简化示例，暂时关闭严格的主机密钥检查
        // 在实际应用中，应根据安全需求谨慎处理主机密钥验证
        java.util.Properties config = new java.util.Properties();
        config.put("StrictHostKeyChecking", "no");
        session.setConfig(config);

        session.connect();

        // 将创建好的连接对象存入Map中，键为host:port的形式
        connectionMap.put(host + ":" + port, session);
        return true;
    }

    // 根据主机名和端口号获取SSH连接对象的方法
    public Session getConnection(String host, int port) {
        return connectionMap.get(host + ":" + port);
    }

    /**
     * 执行SSH命令的方法
     * @param host
     * @param port
     * @param command
     * @return
     * @throws Exception
     */
    @Async
    public String executeCommand(String host, int port, String command) throws Exception {
        Session session = getConnection(host, port);
        if (session == null) {
            // // TODO: 2024/11/21 是否需要重连，从哪里获取用户名密码
            throw new IllegalArgumentException("未找到指定主机和端口的SSH连接，请重连");
        }

        // 读取命令执行结果
        StringBuilder result = new StringBuilder();
        Channel channel = session.openChannel("exec");
        ((ChannelExec) channel).setCommand(command);
        channel.setInputStream(null);
        ((ChannelExec) channel).setErrStream(System.err);

        channel.connect();

        byte[] buffer = new byte[1024];
        InputStream in = channel.getInputStream();
        while (true) {
            while (in.available() > 0) {
                int i = in.read(buffer, 0, 1024);
                if (i < 0) {
                    break;
                }
                result.append(new String(buffer, 0, i));
            }
            if (channel.isClosed()) {
                break;
            }
            try {
                Thread.sleep(1000);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        channel.disconnect();

        return result.toString();
    }

    // channel不关闭，继续执行
    @Async
    public List<String> shellCommand(String host, int port, String command) throws Exception {
        List<String> result = new ArrayList<>();
        Session session = getConnection(host, port);
        if (session == null) {
            throw new IllegalArgumentException("未找到指定主机和端口的SSH连接，请重连");
        }
        // 获取Channel
        /*Channel channel = channelMap.get(host + ":" + port);
        if (channel == null) {
            channel = session.openChannel("shell");
            channel.connect();
            channelMap.put(host + ":" + port, channel);
        }*/
        Channel channel = session.openChannel("shell");
        channel.connect();
//        ((ChannelShell)channel).start();
        InputStream inputStream = channel.getInputStream();
        OutputStream outputStream = channel.getOutputStream();
//        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
        // 发送Telnet命令，假设目标是连接到一个Telnet服务
        outputStream.write(("telnet 127.0.0.1 9000\n").getBytes());
        String username = "router";
        outputStream.write((username + "\n").getBytes());
        String password = "fos";
        outputStream.write((password + "\n").getBytes());
        outputStream.flush();
        outputStream.write((command + "\n").getBytes());
        outputStream.flush();
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
            String line;
            while (reader != null && (line = reader.readLine()) != null) {
                System.out.println(line);
                result.add(line);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        channel.disconnect();
        return result;
    }

    /**
     * 关闭指定主机和端口的SSH连接的方法
     * @param host
     * @param port
     */
    public void closeCurrentConnections(String host, int port) {
        Session session = getConnection(host, port);
        if (session != null) {
            session.disconnect();
        }
        connectionMap.remove(host + ":" + port);
    }


    /**
     * 关闭所有SSH连接的方法
     */
    public void closeAllConnections() {
        for (Session session : connectionMap.values()) {
            session.disconnect();
        }
        connectionMap.clear();
    }
}